'''
Created on Sep 21, 2009

@author: mkiyer
'''

import veggie.coverage.covdb as covdb
from veggie.sample.samplegroup import parse_samplegroups_xml
from optparse import OptionParser
import collections
import glob
import h5py
import logging
import numpy as np
import os
import sys

import matplotlib.pyplot as plt
import matplotlib.cm as cm

class UTR:
    def __init__(self):
        pass
    def __str__(self):
        return '%s\t%d\t%d\t%s\t%s\t%d\t%d' % \
            (self.chrom, self.start, self.end, self.name,
             self.strand, self.polya, self.polya+1)

def read_utr_polya_file(fhd):
    '''
    symbol 
    accession
    chrom
    strand
    utr_tx_start
    utr_tx_end
    utr_cds_start
    utr_cds_end
    polya_count
    polya_starts
    polya_ends
    polya_sites
    '''
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline is None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        u = UTR()
        u.symbol = thisfields[0]
        u.accession = thisfields[1]
        u.chrom = thisfields[2]
        u.strand = thisfields[3]
        u.tx_start = int(thisfields[4])
        u.tx_end = int(thisfields[5])
        u.cds_start = int(thisfields[6])
        u.cds_end = int(thisfields[7])
        u.polya_count = int(thisfields[8])
        u.polya_starts = map(int, thisfields[9].split(','))
        u.polya_ends = map(int, thisfields[10].split(','))
        u.polya_sites = map(int, thisfields[11].split(','))
        yield u 


def calc_utr_polya_coverage(utr, dsets):
    # get coverage of the coding part of the 3'UTR
    cds_cov = covdb.calc_avg_cov_rpkm(utr.chrom, utr.cds_start, utr.cds_end, dsets)
    # for each polya site in the UTR get proximal/distal coverage
    utr_covs = []
    for i in xrange(utr.polya_count):
        start, end, polya = utr.polya_starts[i], utr.polya_ends[i], utr.polya_sites[i]
        proximal = covdb.calc_avg_cov_rpkm(utr.chrom, start, polya, dsets)
        distal = covdb.calc_avg_cov_rpkm(utr.chrom, polya, end, dsets)        
        if utr.strand == '+':
            utr_covs.append((proximal, distal))
        else:
            utr_covs.append((distal, proximal))
    # return the coverage values
    return cds_cov, utr_covs
        

def get_output_header():
    return ('\t'.join(['#symbol',
                       'accession',
                       'chrom',
                       'strand',
                       'cds_exon_start',
                       'cds_exon_end',
                       'exper_cds_cov',
                       'control_cds_cov',
                       'utr_start',
                       'utr_end',
                       'polya_site',
                       'exper_proximal',
                       'exper_distal',
                       'control_proximal',
                       'control_distal']))

def gen_utr_cov_output(utr, e_cds_cov, e_utr_covs, c_cds_cov, c_utr_covs):
    # template for the coverage values of this utr
    template = [u.symbol,
                u.accession,
                u.chrom,
                u.strand,
                str(u.cds_start),
                str(u.cds_end),
                '%.3f' % e_cds_cov,
                '%.3f' % c_cds_cov]
    
    for i in xrange(utr.polya_count):
        thisline = list(template)
        thisutr = map(str, [utr.polya_starts[i],
                            utr.polya_ends[i],
                            utr.polya_sites[i],
                            '%.3f' % e_utr_covs[i][0],
                            '%.3f' % e_utr_covs[i][1],
                            '%.3f' % c_utr_covs[i][0],
                            '%.3f' % c_utr_covs[i][1]])        
        thisline.extend(thisutr)
        yield '\t'.join(thisline)


if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] <polya_table.bed>")    
    optionparser.add_option("-e", "--experiment", dest="exper",
                            help="sample name of experiment")
    optionparser.add_option("-c", "--control", dest="control",
                            help="sample name of control")
    optionparser.add_option("--covdb", dest="covdb",
                            help="path to coverage database")
    optionparser.add_option("-o", "--output", dest="outfile",
                            help="output file [default: %default]")
    (options, args) = optionparser.parse_args()
    if options.outfile == None:
        outfhd = sys.stdout
    else:
        outfhd = open(options.outfile, 'w')
    if options.covdb is None:
        optionparser.error("no HDF5 coverage file specified")
    if len(args) == 0:
        optionparser.error("no polya table file specified in args")

    # configure logging
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(levelname)s %(message)s')
    # open the coverage database file
    h5file = h5py.File(options.covdb, 'r')
    # get treatment and control groups
    e_dsets = covdb.get_datasets_from_sample_name(options.exper, h5file)
    if not e_dsets:
        optionparser.error("No lanes found for experimental sample: %s" % (options.exper))
    c_dsets = covdb.get_datasets_from_sample_name(options.control, h5file)
    if not c_dsets:
        optionparser.error("No lanes found for control sample: %s" % (options.control))    

    # read utr data
    utrs = list(read_utr_polya_file(open(args[0])))
    logging.debug("read %d 3'UTRs" % len(utrs))
    
    num_utrs = 0
    polya_count = 0
    pmax, pmin = 0, 100
    for u in utrs:
        if u.polya_count != 0:
            num_utrs += 1
            polya_count += u.polya_count
            pmax = max(pmax, u.polya_count)
            pmin = min(pmin, u.polya_count)
    logging.debug("average polyA sites per protein-coding UTR=%f" % (float(polya_count)/num_utrs))
    logging.debug("maximum polyA sites found in a single UTR=%d", pmax)
        
    logging.debug('Finding coverage proximal/distal to each polyA site...')
    lines_finished = 0    
    outfhd.write('%s\n' % get_output_header())
    # start the detection
    for u in utrs:
        lines_finished += 1
        if lines_finished % 1000 == 0:
            logging.debug("Finished %d / %d UTRs" % (lines_finished, len(utrs)))

        # only consider UTRs with polya sites
        if u.polya_count == 0:
            continue
        # TODO: only consider UTRs of protein coding genes for now
        if u.cds_start == u.cds_end:
            continue
        e_cds_cov, e_utr_covs = calc_utr_polya_coverage(u, e_dsets)
        c_cds_cov, c_utr_covs = calc_utr_polya_coverage(u, c_dsets)

        for line in gen_utr_cov_output(u, e_cds_cov, e_utr_covs, c_cds_cov, c_utr_covs):
            outfhd.write('%s\n' % line)

    # done
    h5file.close()
    logging.shutdown()